# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

from . import test_data_movement_examples as data_movement
from . import test_core_examples as core
from . import test_tensor_creation_examples as tensor_creation
from . import test_matrix_multiplication_examples as matrix_multiplication
from . import test_eltwise_unary_examples as eltwise_unary
from . import test_complex_unary_examples as complex_unary
from . import test_backward_examples as backward
from . import test_pointwise_binary_examples as pointwise_binary
from . import test_pointwise_ternary_examples as pointwise_ternary
from . import test_losses_examples as losses
from . import test_reduction_examples as reduction
from . import test_normalization_examples as normalization
from . import test_embedding_examples as embedding
from . import test_pooling_examples as pooling

FUNCTION_TO_EXAMPLES_MAPPING_DICT = {
    # Core
    "ttnn.as_tensor": core.test_as_tensor,
    "ttnn.from_torch": core.test_from_torch,
    "ttnn.to_torch": core.test_to_torch,
    "ttnn.to_device": core.test_to_device,
    "ttnn.from_device": core.test_from_device,
    "ttnn.to_layout": core.test_to_layout,
    "ttnn.dump_tensor": core.test_dump_tensor,
    "ttnn.load_tensor": core.test_load_tensor,
    "ttnn.deallocate": core.test_deallocate,
    "ttnn.reallocate": core.test_reallocate,
    "ttnn.to_memory_config": core.test_to_memory_config,
    # Tensor Creation
    "ttnn.arange": tensor_creation.test_arrange,
    "ttnn.empty": tensor_creation.test_empty,
    "ttnn.empty_like": tensor_creation.test_empty_like,
    "ttnn.zeros": tensor_creation.test_zeros,
    "ttnn.zeros_like": tensor_creation.test_zeros_like,
    "ttnn.ones": tensor_creation.test_ones,
    "ttnn.ones_like": tensor_creation.test_ones_like,
    "ttnn.full": tensor_creation.test_full,
    "ttnn.full_like": tensor_creation.test_full_like,
    "ttnn.rand": tensor_creation.test_rand,
    "ttnn.from_buffer": tensor_creation.test_from_buffer,
    # Matrix Multiplication
    "ttnn.matmul": matrix_multiplication.test_matmul,
    "ttnn.linear": matrix_multiplication.test_linear,
    "ttnn.addmm": matrix_multiplication.test_addmm,
    "ttnn.sparse_matmul": matrix_multiplication.test_sparse_matmul,
    # Pointwise Unary
    "ttnn.abs": eltwise_unary.test_abs,
    "ttnn.acos": eltwise_unary.test_acos,
    "ttnn.acosh": eltwise_unary.test_acosh,
    "ttnn.asin": eltwise_unary.test_asin,
    "ttnn.asinh": eltwise_unary.test_asinh,
    "ttnn.atan": eltwise_unary.test_atan,
    "ttnn.atanh": eltwise_unary.test_atanh,
    "ttnn.bitcast": eltwise_unary.test_bitcast,
    "ttnn.bitwise_not": eltwise_unary.test_bitwise_not,
    "ttnn.bitwise_left_shift": eltwise_unary.test_bitwise_left_shift,
    "ttnn.bitwise_right_shift": eltwise_unary.test_bitwise_right_shift,
    "ttnn.cbrt": eltwise_unary.test_cbrt,
    "ttnn.ceil": eltwise_unary.test_ceil,
    "ttnn.celu": eltwise_unary.test_celu,
    "ttnn.clamp": eltwise_unary.test_clamp,
    "ttnn.clip": eltwise_unary.test_clip,
    "ttnn.clone": eltwise_unary.test_clone,
    "ttnn.cos": eltwise_unary.test_cos,
    "ttnn.cosh": eltwise_unary.test_cosh,
    "ttnn.deg2rad": eltwise_unary.test_deg2rad,
    "ttnn.digamma": eltwise_unary.test_digamma,
    "ttnn.elu": eltwise_unary.test_elu,
    "ttnn.eqz": eltwise_unary.test_eqz,
    "ttnn.erf": eltwise_unary.test_erf,
    "ttnn.erfc": eltwise_unary.test_erfc,
    "ttnn.erfinv": eltwise_unary.test_erfinv,
    "ttnn.exp": eltwise_unary.test_exp,
    "ttnn.exp2": eltwise_unary.test_exp2,
    "ttnn.expm1": eltwise_unary.test_expm1,
    "ttnn.fill": eltwise_unary.test_fill,
    "ttnn.floor": eltwise_unary.test_floor,
    "ttnn.frac": eltwise_unary.test_frac,
    "ttnn.geglu": eltwise_unary.test_geglu,
    "ttnn.gelu": eltwise_unary.test_gelu,
    "ttnn.glu": eltwise_unary.test_glu,
    "ttnn.gez": eltwise_unary.test_gez,
    "ttnn.gtz": eltwise_unary.test_gtz,
    "ttnn.hardshrink": eltwise_unary.test_hardshrink,
    "ttnn.hardsigmoid": eltwise_unary.test_hardsigmoid,
    "ttnn.hardswish": eltwise_unary.test_hardswish,
    "ttnn.hardtanh": eltwise_unary.test_hardtanh,
    "ttnn.heaviside": eltwise_unary.test_heaviside,
    "ttnn.i0": eltwise_unary.test_i0,
    "ttnn.identity": eltwise_unary.test_identity,
    "ttnn.isfinite": eltwise_unary.test_isfinite,
    "ttnn.isinf": eltwise_unary.test_isinf,
    "ttnn.isnan": eltwise_unary.test_isnan,
    "ttnn.isneginf": eltwise_unary.test_isneginf,
    "ttnn.isposinf": eltwise_unary.test_isposinf,
    "ttnn.leaky_relu": eltwise_unary.test_leaky_relu,
    "ttnn.lez": eltwise_unary.test_lez,
    "ttnn.lgamma": eltwise_unary.test_lgamma,
    "ttnn.log": eltwise_unary.test_log,
    "ttnn.log10": eltwise_unary.test_log10,
    "ttnn.log1p": eltwise_unary.test_log1p,
    "ttnn.log2": eltwise_unary.test_log2,
    "ttnn.log_sigmoid": eltwise_unary.test_log_sigmoid,
    "ttnn.logical_not": eltwise_unary.test_logical_not,
    "ttnn.logical_not_": eltwise_unary.test_logical_not_,
    "ttnn.logit": eltwise_unary.test_logit,
    "ttnn.ltz": eltwise_unary.test_ltz,
    "ttnn.mish": eltwise_unary.test_mish,
    "ttnn.multigammaln": eltwise_unary.test_multigammaln,
    "ttnn.neg": eltwise_unary.test_neg,
    "ttnn.nez": eltwise_unary.test_nez,
    "ttnn.normalize_global": eltwise_unary.test_normalize_global,
    "ttnn.normalize_hw": eltwise_unary.test_normalize_hw,
    "ttnn.polygamma": eltwise_unary.test_polygamma,
    "ttnn.prelu": eltwise_unary.test_prelu,
    "ttnn.rad2deg": eltwise_unary.test_rad2deg,
    "ttnn.rdiv": eltwise_unary.test_rdiv,
    "ttnn.reciprocal": eltwise_unary.test_reciprocal,
    "ttnn.reglu": eltwise_unary.test_reglu,
    "ttnn.relu": eltwise_unary.test_relu,
    "ttnn.relu_max": eltwise_unary.test_relu_max,
    "ttnn.relu_min": eltwise_unary.test_relu_min,
    "ttnn.relu6": eltwise_unary.test_relu6,
    "ttnn.remainder": eltwise_unary.test_remainder,
    "ttnn.round": eltwise_unary.test_round,
    "ttnn.rpow": eltwise_unary.test_rpow,
    "ttnn.rsqrt": eltwise_unary.test_rsqrt,
    "ttnn.selu": eltwise_unary.test_selu,
    "ttnn.sigmoid": eltwise_unary.test_sigmoid,
    "ttnn.sigmoid_accurate": eltwise_unary.test_sigmoid_accurate,
    "ttnn.sign": eltwise_unary.test_sign,
    "ttnn.signbit": eltwise_unary.test_signbit,
    "ttnn.silu": eltwise_unary.test_silu,
    "ttnn.sin": eltwise_unary.test_sin,
    "ttnn.sinh": eltwise_unary.test_sinh,
    "ttnn.softplus": eltwise_unary.test_softplus,
    "ttnn.softshrink": eltwise_unary.test_softshrink,
    "ttnn.softsign": eltwise_unary.test_softsign,
    "ttnn.sqrt": eltwise_unary.test_sqrt,
    "ttnn.square": eltwise_unary.test_square,
    "ttnn.swiglu": eltwise_unary.test_swiglu,
    "ttnn.swish": eltwise_unary.test_swish,
    "ttnn.tan": eltwise_unary.test_tan,
    "ttnn.tanh": eltwise_unary.test_tanh,
    "ttnn.tanhshrink": eltwise_unary.test_tanhshrink,
    "ttnn.threshold": eltwise_unary.test_threshold,
    "ttnn.tril": eltwise_unary.test_tril,
    "ttnn.triu": eltwise_unary.test_triu,
    "ttnn.trunc": eltwise_unary.test_trunc,
    "ttnn.unary_chain": eltwise_unary.test_unary_chain,
    # Complex Unary
    "ttnn.real": complex_unary.test_real,
    "ttnn.imag": complex_unary.test_imag,
    "ttnn.angle": complex_unary.test_angle,
    "ttnn.is_imag": complex_unary.test_is_imag,
    "ttnn.is_real": complex_unary.test_is_real,
    "ttnn.conj": complex_unary.test_conj,
    "ttnn.polar": complex_unary.test_polar,
    "ttnn.alt_complex_rotate90": complex_unary.test_alt_complex_rotate90,
    # Pointwise Unary Backward Operations
    "ttnn.clamp_bw": backward.test_clamp_bw,
    "ttnn.clip_bw": backward.test_clip_bw,
    "ttnn.hardtanh_bw": backward.test_hardtanh_bw,
    "ttnn.threshold_bw": backward.test_threshold_bw,
    "ttnn.softplus_bw": backward.test_softplus_bw,
    "ttnn.rdiv_bw": backward.test_rdiv_bw,
    "ttnn.pow_bw": backward.test_pow_bw,
    "ttnn.exp_bw": backward.test_exp_bw,
    "ttnn.tanh_bw": backward.test_tanh_bw,
    "ttnn.sqrt_bw": backward.test_sqrt_bw,
    "ttnn.multigammaln_bw": backward.test_multigammaln_bw,
    "ttnn.lgamma_bw": backward.test_lgamma_bw,
    "ttnn.fill_bw": backward.test_fill_bw,
    "ttnn.hardsigmoid_bw": backward.test_hardsigmoid_bw,
    "ttnn.cos_bw": backward.test_cos_bw,
    "ttnn.acosh_bw": backward.test_acosh_bw,
    "ttnn.acos_bw": backward.test_acos_bw,
    "ttnn.atan_bw": backward.test_atan_bw,
    "ttnn.rad2deg_bw": backward.test_rad2deg_bw,
    "ttnn.frac_bw": backward.test_frac_bw,
    "ttnn.trunc_bw": backward.test_trunc_bw,
    "ttnn.log_sigmoid_bw": backward.test_log_sigmoid_bw,
    "ttnn.fill_zero_bw": backward.test_fill_zero_bw,
    "ttnn.i0_bw": backward.test_i0_bw,
    "ttnn.tan_bw": backward.test_tan_bw,
    "ttnn.sigmoid_bw": backward.test_sigmoid_bw,
    "ttnn.rsqrt_bw": backward.test_rsqrt_bw,
    "ttnn.neg_bw": backward.test_neg_bw,
    "ttnn.relu_bw": backward.test_relu_bw,
    "ttnn.logit_bw": backward.test_logit_bw,
    "ttnn.hardshrink_bw": backward.test_hardshrink_bw,
    "ttnn.softshrink_bw": backward.test_softshrink_bw,
    "ttnn.leaky_relu_bw": backward.test_leaky_relu_bw,
    "ttnn.elu_bw": backward.test_elu_bw,
    "ttnn.celu_bw": backward.test_celu_bw,
    "ttnn.rpow_bw": backward.test_rpow_bw,
    "ttnn.floor_bw": backward.test_floor_bw,
    "ttnn.round_bw": backward.test_round_bw,
    "ttnn.log_bw": backward.test_log_bw,
    "ttnn.relu6_bw": backward.test_relu6_bw,
    "ttnn.abs_bw": backward.test_abs_bw,
    "ttnn.silu_bw": backward.test_silu_bw,
    "ttnn.selu_bw": backward.test_selu_bw,
    "ttnn.square_bw": backward.test_square_bw,
    "ttnn.prod_bw": backward.test_prod_bw,
    "ttnn.hardswish_bw": backward.test_hardswish_bw,
    "ttnn.tanhshrink_bw": backward.test_tanhshrink_bw,
    "ttnn.atanh_bw": backward.test_atanh_bw,
    "ttnn.asin_bw": backward.test_asin_bw,
    "ttnn.asinh_bw": backward.test_asinh_bw,
    "ttnn.sin_bw": backward.test_sin_bw,
    "ttnn.sinh_bw": backward.test_sinh_bw,
    "ttnn.log10_bw": backward.test_log10_bw,
    "ttnn.log1p_bw": backward.test_log1p_bw,
    "ttnn.erfc_bw": backward.test_erfc_bw,
    "ttnn.ceil_bw": backward.test_ceil_bw,
    "ttnn.softsign_bw": backward.test_softsign_bw,
    "ttnn.cosh_bw": backward.test_cosh_bw,
    "ttnn.logiteps_bw": backward.test_logiteps_bw,
    "ttnn.log2_bw": backward.test_log2_bw,
    "ttnn.sign_bw": backward.test_sign_bw,
    "ttnn.div_no_nan_bw": backward.test_div_no_nan_bw,
    "ttnn.exp2_bw": backward.test_exp2_bw,
    "ttnn.expm1_bw": backward.test_expm1_bw,
    "ttnn.reciprocal_bw": backward.test_reciprocal_bw,
    "ttnn.digamma_bw": backward.test_digamma_bw,
    "ttnn.erfinv_bw": backward.test_erfinv_bw,
    "ttnn.erf_bw": backward.test_erf_bw,
    "ttnn.deg2rad_bw": backward.test_deg2rad_bw,
    "ttnn.polygamma_bw": backward.test_polygamma_bw,
    "ttnn.gelu_bw": backward.test_gelu_bw,
    "ttnn.repeat_bw": backward.test_repeat_bw,
    "ttnn.polar_bw": backward.test_polar_bw,
    "ttnn.imag_bw": backward.test_imag_bw,
    "ttnn.real_bw": backward.test_real_bw,
    "ttnn.angle_bw": backward.test_angle_bw,
    "ttnn.conj_bw": backward.test_conj_bw,
    # Pointwise Binary
    "ttnn.add": pointwise_binary.test_add,
    "ttnn.addalpha": pointwise_binary.test_addalpha,
    "ttnn.subalpha": pointwise_binary.test_subalpha,
    "ttnn.multiply": pointwise_binary.test_multiply,
    "ttnn.subtract": pointwise_binary.test_subtract,
    "ttnn.div": pointwise_binary.test_div,
    "ttnn.div_no_nan": pointwise_binary.test_div_no_nan,
    "ttnn.floor_div": pointwise_binary.test_floor_div,
    "ttnn.remainder": pointwise_binary.test_remainder,
    "ttnn.fmod": pointwise_binary.test_fmod,
    "ttnn.gcd": pointwise_binary.test_gcd,
    "ttnn.lcm": pointwise_binary.test_lcm,
    "ttnn.logical_and_": pointwise_binary.test_logical_and_,
    "ttnn.logical_or_": pointwise_binary.test_logical_or_,
    "ttnn.logical_xor_": pointwise_binary.test_logical_xor_,
    "ttnn.rpow": pointwise_binary.test_rpow,
    "ttnn.rsub": pointwise_binary.test_rsub,
    "ttnn.ldexp": pointwise_binary.test_ldexp,
    "ttnn.logical_and": pointwise_binary.test_logical_and,
    "ttnn.logical_or": pointwise_binary.test_logical_or,
    "ttnn.logical_xor": pointwise_binary.test_logical_xor,
    "ttnn.bitwise_and": pointwise_binary.test_bitwise_and,
    "ttnn.bitwise_or": pointwise_binary.test_bitwise_or,
    "ttnn.bitwise_xor": pointwise_binary.test_bitwise_xor,
    "ttnn.logaddexp": pointwise_binary.test_logaddexp,
    "ttnn.logaddexp2": pointwise_binary.test_logaddexp2,
    "ttnn.hypot": pointwise_binary.test_hypot,
    "ttnn.xlogy": pointwise_binary.test_xlogy,
    "ttnn.squared_difference": pointwise_binary.test_squared_difference,
    "ttnn.gt": pointwise_binary.test_gt,
    "ttnn.gt_": pointwise_binary.test_gt_,
    "ttnn.lt_": pointwise_binary.test_lt_,
    "ttnn.ge_": pointwise_binary.test_ge_,
    "ttnn.le_": pointwise_binary.test_le_,
    "ttnn.eq_": pointwise_binary.test_eq_,
    "ttnn.ne_": pointwise_binary.test_ne_,
    "ttnn.ge": pointwise_binary.test_ge,
    "ttnn.lt": pointwise_binary.test_lt,
    "ttnn.le": pointwise_binary.test_le,
    "ttnn.eq": pointwise_binary.test_eq,
    "ttnn.ne": pointwise_binary.test_ne,
    "ttnn.isclose": pointwise_binary.test_isclose,
    "ttnn.nextafter": pointwise_binary.test_nextafter,
    "ttnn.maximum": pointwise_binary.test_maximum,
    "ttnn.minimum": pointwise_binary.test_minimum,
    "ttnn.outer": pointwise_binary.test_outer,
    "ttnn.pow": pointwise_binary.test_pow,
    "ttnn.polyval": pointwise_binary.test_polyval,
    "ttnn.scatter": pointwise_binary.test_scatter,
    "ttnn.atan2": pointwise_binary.test_atan2,
    # Pointwise Binary Backward
    "ttnn.add_bw": pointwise_binary.test_add_bw,
    "ttnn.assign_bw": pointwise_binary.test_assign_bw,
    "ttnn.atan2_bw": pointwise_binary.test_atan2_bw,
    "ttnn.bias_gelu_bw": pointwise_binary.test_bias_gelu_bw,
    "ttnn.div_bw": pointwise_binary.test_div_bw,
    "ttnn.embedding_bw": pointwise_binary.test_embedding_bw,
    "ttnn.fmod_bw": pointwise_binary.test_fmod_bw,
    "ttnn.remainder_bw": pointwise_binary.test_remainder_bw,
    "ttnn.addalpha_bw": pointwise_binary.test_addalpha_bw,
    "ttnn.subalpha_bw": pointwise_binary.test_subalpha_bw,
    "ttnn.xlogy_bw": pointwise_binary.test_xlogy_bw,
    "ttnn.hypot_bw": pointwise_binary.test_hypot_bw,
    "ttnn.ldexp_bw": pointwise_binary.test_ldexp_bw,
    "ttnn.logaddexp_bw": pointwise_binary.test_logaddexp_bw,
    "ttnn.logaddexp2_bw": pointwise_binary.test_logaddexp2_bw,
    "ttnn.mul_bw": pointwise_binary.test_mul_bw,
    "ttnn.sub_bw": pointwise_binary.test_sub_bw,
    "ttnn.squared_difference_bw": pointwise_binary.test_squared_difference_bw,
    "ttnn.concat_bw": pointwise_binary.test_concat_bw,
    "ttnn.rsub_bw": pointwise_binary.test_rsub_bw,
    "ttnn.min_bw": pointwise_binary.test_min_bw,
    "ttnn.max_bw": pointwise_binary.test_max_bw,
    # Pointwise Ternary
    "ttnn.addcdiv": pointwise_ternary.test_addcdiv,
    "ttnn.addcmul": pointwise_ternary.test_addcmul,
    "ttnn.mac": pointwise_ternary.test_mac,
    "ttnn.where": pointwise_ternary.test_where,
    "ttnn.lerp": pointwise_ternary.test_lerp,
    # Pointwise Ternary Backward
    "ttnn.addcmul_bw": pointwise_ternary.test_addcmul_bw,
    "ttnn.addcdiv_bw": pointwise_ternary.test_addcdiv_bw,
    "ttnn.where_bw": pointwise_ternary.test_where_bw,
    "ttnn.lerp_bw": pointwise_ternary.test_lerp_bw,
    # Losses
    "ttnn.l1_loss": losses.test_l1_loss,
    "ttnn.mse_loss": losses.test_mse_loss,
    # Reduction
    "ttnn.cumprod": reduction.test_cumprod,
    "ttnn.max": reduction.test_max,
    "ttnn.mean": reduction.test_mean,
    "ttnn.min": reduction.test_min,
    "ttnn.std": reduction.test_std,
    "ttnn.sum": reduction.test_sum,
    "ttnn.var": reduction.test_var,
    "ttnn.argmax": reduction.test_argmax,
    "ttnn.prod": reduction.test_prod,
    "ttnn.topk": reduction.test_topk,
    "ttnn.cumsum": reduction.test_cumsum,
    "ttnn.ema": reduction.test_ema,
    "ttnn.moe": reduction.test_moe,
    "ttnn.manual_seed": reduction.test_manual_seed,
    # Data movement
    "ttnn.concat": data_movement.test_concat,
    "ttnn.nonzero": data_movement.test_nonzero,
    "ttnn.pad": data_movement.test_pad,
    "ttnn.permute": data_movement.test_permute,
    "ttnn.reshape": data_movement.test_reshape,
    "ttnn.repeat": data_movement.test_repeat,
    "ttnn.repeat_interleave": data_movement.test_repeat_interleave,
    "ttnn.slice": data_movement.test_slice,
    "ttnn.tilize": data_movement.test_tilize,
    "ttnn.tilize_with_val_padding": data_movement.test_tilize_with_val_padding,
    "ttnn.fill_rm": data_movement.test_fill_rm,
    "ttnn.fill_ones_rm": data_movement.test_fill_ones_rm,
    "ttnn.untilize": data_movement.test_untilize,
    "ttnn.untilize_with_unpadding": data_movement.test_untilize_with_unpadding,
    "ttnn.indexed_fill": data_movement.test_indexed_fill,
    "ttnn.gather": data_movement.test_gather,
    "ttnn.sort": data_movement.test_sort,
    # Normalization
    "ttnn.group_norm": normalization.test_group_norm,
    "ttnn.layer_norm": normalization.test_layer_norm,
    "ttnn.layer_norm_pre_all_gather": normalization.test_layernorm_distributed,
    "ttnn.layer_norm_post_all_gather": normalization.test_layernorm_distributed,
    "ttnn.rms_norm": normalization.test_rms_norm,
    "ttnn.rms_norm_pre_all_gather": normalization.test_rms_norm_distributed,
    "ttnn.rms_norm_post_all_gather": normalization.test_rms_norm_distributed,
    "ttnn.batch_norm": normalization.test_batch_norm,
    "ttnn.softmax": normalization.test_softmax,
    "ttnn.scale_mask_softmax": normalization.test_scale_mask_softmax,
    "ttnn.softmax_in_place": normalization.test_softmax_in_place,
    "ttnn.scale_mask_softmax_in_place": normalization.test_scale_mask_softmax_in_place,
    "ttnn.scale_causal_mask_hw_dims_softmax_in_place": normalization.test_scale_causal_mask_hw_dims_softmax_in_place,
    # Normalization Program Configs
    "ttnn.SoftmaxDefaultProgramConfig": normalization.test_softmax_default_program_config,
    "ttnn.SoftmaxShardedMultiCoreProgramConfig": normalization.test_scale_mask_softmax_in_place,
    # Transformers
    # "ttnn.transformer.split_query_key_value_and_split_heads": transformer.test_split_query_key_value_and_split_heads, # Lack of example
    # "ttnn.transformer.concatenate_heads": transformer.test_concatenate_heads, # Lack of example
    # "ttnn.transformer.attention_softmax": transformer.test_attention_softmax, # Lack of example
    # "ttnn.transformer.attention_softmax_": transformer.test_attention_softmax_, # Lack of example
    # "ttnn.experimental.rotary_embedding": transformer.test_rotary_embedding, # Lack of example
    # "ttnn.transformer.scaled_dot_product_attention": transformer.test_scaled_dot_product_attention, # Lack of example
    # "ttnn.transformer.scaled_dot_product_attention_decode": transformer.test_scaled_dot_product_attention_decode, # Lack of example
    # CCL
    # "ttnn.all_gather": ccl.test_all_gather, # Non-working example
    # "ttnn.reduce_scatter": ccl.test_reduce_scatter, # Non-working example
    # "ttnn.all_reduce": ccl.test_all_reduce, # Non-working example
    # Embedding
    "ttnn.embedding": embedding.test_embedding,
    # Convolution
    # "ttnn.conv1d": convolution.test_conv1d, # Lack of example
    # "ttnn.conv2d": convolution.test_conv2d, # Lack of example
    # "ttnn.experimental.conv3d": convolution.test_conv3d, # Lack of example
    # "ttnn.conv_transpose2d": convolution.test_conv_transpose2d, # Lack of example
    # "ttnn.prepare_conv_weights": convolution.test_prepare_conv_weights, # Lack of example
    # "ttnn.prepare_conv_bias": convolution.test_prepare_conv_bias, # Lack of example
    # "ttnn.prepare_conv_transpose2d_weights": convolution.test_prepare_conv_transpose2d_weights, # Lack of example
    # "ttnn.prepare_conv_transpose2d_bias": convolution.test_prepare_conv_transpose2d_bias, # Lack of example
    # Pooling
    "ttnn.global_avg_pool2d": pooling.test_global_avg_pool2d,
    "ttnn.max_pool2d": pooling.test_max_pool2d,
    "ttnn.avg_pool2d": pooling.test_avg_pool2d,
    # Vision
    # "ttnn.upsample": vision.test_upsample, # Lack of example
    # KV Cache
    # "ttnn.kv_cache.fill_cache_for_user_":kv_cache.test_fill_cache_for_user_, # Lack of example
    # "ttnn.kv_cache.update_cache_for_token_":kv_cache.test_update_cache_for_token_, # Lack of example
}
